import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
from zhipuai import ZhipuAI
import json
from streamlit_echarts import st_echarts

# 初始化会话状态
# 在文件开头移除history初始化
# 删除这行代码：
# if 'history' not in st.session_state:
#     st.session_state.history = []

# 初始化智谱AI客户端
def init_ai(api_key):
    if not api_key or len(api_key) < 10:  # 简单验证API Key长度
        raise ValueError("无效的API Key")
    return ZhipuAI(api_key=api_key)

# 加载招聘数据
@st.cache_data
def load_data():
    data = pd.read_csv('拉勾网2023招聘数据.csv')
    
    def parse_salary(salary_str):
        if '-' in salary_str:
            low, high = salary_str.lower().replace('k', '').split('-')
            return (float(low) + float(high)) / 2 * 1000
        return None
    
    data['salary_mid'] = data['salary'].apply(parse_salary)
    return data

# 新增工具函数（修改query_job_demand描述并补充薪资最高工具）
tools = [
    {
        "type": "function",
        "function": {
            "name": "query_language_jobs",
            "description": "查询指定编程语言相关岗位的详细信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "language": {"type": "string", "description": "编程语言名称"},
                    "top_n": {"type": "number", "description": "返回前N个岗位（可选，默认10）"}
                },
                "required": ["language"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_average_salary",
            "description": "查询某城市某职位的平均薪资",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "job_name": {"type": "string", "description": "职位名称"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_salary_range",
            "description": "查询某城市某职位的薪资区间",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "job_name": {"type": "string", "description": "职位名称"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_job_demand",
            "description": "查询某城市某职位的岗位需求量（岗位数量）",  # 修正描述
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "job_name": {"type": "string", "description": "职位名称"}
                },
                "required": ["city", "job_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_max_salary_job",
            "description": "查询某城市薪资最高的岗位",  # 新增薪资最高岗位工具
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"},
                    "top_n": {"type": "number", "description": "返回前N个薪资最高岗位（可选，默认5）"}
                },
                "required": ["city"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_language_related_jobs",  # 重命名工具函数
            "description": "查询掌握指定编程语言可以从事的岗位及其需求量（如Python、C++等）",
            "parameters": {
                "type": "object",
                "properties": {
                    "language": {"type": "string", "description": "编程语言名称"},
                    "top_n": {"type": "number", "description": "返回前N个岗位（可选，默认5）"}
                },
                "required": ["language"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "query_cpp_learning_recommendations",
            "description": "查询C++相关的技术学习推荐（如推荐技术栈、学习路径等）",
            "parameters": {
                "type": "object",
                "properties": {
                    "top_n": {"type": "number", "description": "返回前N个推荐（可选，默认5）"}
                }
            }
        }
    }
]

# 本地函数实现
def query_language_jobs(data, language="Java", top_n=10):
    """查询指定编程语言相关岗位的详细信息"""
    try:
        # 新增有效性校验
        required_columns = ['positionName', 'companyFullName', 'city', 'salary', 'workYear', 'education', 'skillLables']
        if not all(col in data.columns for col in required_columns):
            return pd.DataFrame()  # 返回空DataFrame避免后续错误

        filtered = data[
            data['skillLables'].str.contains(language, case=False, na=False)
        ].sort_values('salary_mid', ascending=False)
        
        # 提取并重命名列
        rdata = filtered[[
            'positionName', 
            'companyFullName', 
            'city', 
            'salary',
            'workYear',
            'education'
        ]].head(top_n)
        
        # 重命名列名为中文
        rdata = rdata.rename(columns={
            'positionName': '职位名称',
            'companyFullName': '公司名称',
            'city': '城市',
            'salary': '薪资范围',
            'workYear': '工作经验',
            'education': '学历要求'
        })
        
        # 处理空值和格式
        rdata.fillna('不限', inplace=True)
        rdata['薪资范围'] = rdata['薪资范围'].str.replace('k', 'K', case=False)
        return rdata
    
    except Exception as e:
        return pd.DataFrame()
        
def query_average_salary(city, job_name, data):
    # 检查必要字段是否存在
    required_columns = ['city', 'positionName', 'salary_mid']
    if not all(col in data.columns for col in required_columns):
        return f"数据缺失必要字段：{[col for col in required_columns if col not in data.columns]}"
    
    # 筛选数据
    job_data = data[(data['city'] == city) & (data['positionName'] == job_name)]
    if job_data.empty:
        return {"error": f"未找到{city}的{job_name}岗位数据"}
    if job_data['salary_mid'].isna().all():
        return {"error": f"{city}的{job_name}岗位薪资数据缺失"}
    
    avg_salary = job_data['salary_mid'].mean()
    return {f"{city}{job_name}平均薪资": avg_salary}  # 返回字典格式

def query_salary_range(city, job_name, data):
    job_data = data[(data['city'] == city) & (data['positionName'] == job_name)]
    if not job_data.empty:
        return (job_data['salary_mid'].min(), job_data['salary_mid'].max())
    return None

def query_job_demand(city, job_name, data):
    job_data = data[(data['city'] == city) & (data['positionName'] == job_name)]
    count = len(job_data)
    return {f"{city}{job_name}岗位数量": count}  # 确保返回字典格式

# 本地函数实现（修改query_top_jobs参数顺序）
def query_top_jobs(city, data, top_n=5):  
    city_data = data[data['city'] == city]
    if len(city_data) == 0:
        return {}
    top_jobs = city_data['positionName'].value_counts().head(top_n)
    return top_jobs.to_dict()

# 新增本地函数：查询某城市薪资最高的岗位（仅返回最高1个）
def query_max_salary_job(city, data, top_n=5):
    city_data = data[data['city'] == city].dropna(subset=['salary_mid'])
    if city_data.empty:
        return {}
    # 按薪资中位数降序排序，取第1个岗位
    top_jobs = city_data.groupby('positionName')['salary_mid'].mean() \
                       .sort_values(ascending=False) \
                       .head(1)  # 强制取第1个
    return top_jobs.to_dict()

# 可视化函数
# 可视化函数（支持柱状图/饼图）
def plot_top_jobs(top_jobs, chart_type='bar', title=None):  # 强制title为None
    if not top_jobs:
        return
    # 新增环形图配置
    if chart_type == 'ring':
        options = {
            "tooltip": {
                "trigger": "item",
                "formatter": "{b}: {c}（推荐指数）"
            },
            "legend": {
                "orient": "vertical",
                "left": "left",
                "data": list(top_jobs.keys())
            },
            "series": [{
                "name": "推荐指数",
                "type": "pie",
                "radius": ['30%', '70%'],  # 关键修改：内外半径形成环形
                "label": {
                    "show": True,
                    "formatter": "{b}: {d}%"
                },
                "data": [{"value": v, "name": k} for k, v in top_jobs.items()],
                "itemStyle": {
                    "borderRadius": 5,  # 新增圆角效果
                    "borderColor": '#fff',
                    "borderWidth": 2
                }
            }]
        }
    elif chart_type == 'pie':
        options = {
            "tooltip": {
                "trigger": "item",
                "formatter": "{b}: {c}（推荐指数）"  # 关键修改：匹配学习推荐场景的提示语
            },
            "legend": {"orient": "vertical", "left": "left", "data": list(top_jobs.keys())},
            "series": [{
                "name": "推荐指数",  # 修改系列名称
                "type": "pie",
                "radius": "70%",
                "data": [{"value": v, "name": k} for k, v in top_jobs.items()],
                "emphasis": {
                    "itemStyle": {"shadowBlur": 10, "shadowOffsetX": 0, "shadowColor": "rgba(0, 0, 0, 0.5)"}
                }
            }]
        }
    else:
        options = {
            "title": {"text": title, "left": "center"},  # 使用动态标题
            "tooltip": {"trigger": "axis", "axisPointer": {"type": "shadow"}},
            "xAxis": {
                "type": "category",
                "data": list(top_jobs.keys()),
                "axisLabel": {"rotate": 45}
            },
            "yAxis": {"type": "value", "name": "岗位数量"},
            "series": [{
                "name": "岗位需求量",
                "type": "bar",
                "data": list(top_jobs.values()),
                "itemStyle": {"color": "#5470C6"}
            }]
        }
    
    st_echarts(
        options=options,
        height=400,
        key=f"{chart_type}_top_jobs_{len(st.session_state.messages)}"
    )

# 处理用户问题函数（关键修改）
# 新增本地函数：查询C++相关岗位
def query_language_related_jobs(data, language="C++", top_n=5):
    # 动态筛选指定语言的岗位
    jobs = data[data['skillLables'].str.contains(language, na=False, case=False)]
    if jobs.empty:
        return {}
    return jobs['positionName'].value_counts().head(top_n).to_dict()
    
    # 统计岗位名称出现次数（取前top_n）
    job_counts = cpp_jobs['positionName'].value_counts().head(top_n)
    return job_counts.to_dict()

def query_cpp_learning_recommendations(data, top_n=5):
    # 筛选包含C++技能的岗位（假设数据字段为'skillLables'，用逗号分隔技能）
    cpp_jobs = data[data['skillLables'].str.contains('C++', na=False, case=False)]
    
    if cpp_jobs.empty:
        # 无实际数据时使用示例推荐
        learning_recommendations = {
            "现代C++特性（C++11/14/17/20）": 9,
            "操作系统与系统编程（Linux内核）": 8,
            "高性能计算与并发编程": 7,
            "C++标准库与STL深度": 6,
            "开源项目贡献（如LLVM）": 5
        }
        # 修复括号未闭合问题
        return dict(sorted(learning_recommendations.items(), key=lambda x: x[1], reverse=True)[:top_n])
    
    # 收集所有关联技能（排除C++本身）
    all_related_skills = []
    for skills in cpp_jobs['skillLables']:
        if isinstance(skills, str):
            skills_list = [s.strip() for s in skills.split(',') if s.strip() != 'C++']
            all_related_skills.extend(skills_list)
    
    # 统计技能出现频率（作为推荐指数）
    skill_counts = pd.Series(all_related_skills).value_counts().head(top_n)
    
    # 转换为{技能名称: 推荐指数}格式
    return skill_counts.to_dict()

# 处理用户问题
def process_user_question(question, data, api_key):
    try:
        client = init_ai(api_key)
        messages = [{"role": "user", "content": question}]
        
        if st.session_state.messages:
            messages.extend([msg for msg in st.session_state.messages[-5:]])
        
        response = client.chat.completions.create(
            model="glm-4",
            messages=messages,
            tools=tools,
            tool_choice="auto"
        )
            
        if response.choices[0].message.tool_calls:
            tool_call = response.choices[0].message.tool_calls[0]
            func_name = tool_call.function.name
            params = json.loads(tool_call.function.arguments)
            
            result = None
            chart_type = None  # Track chart type
            title = None  
            if func_name == "query_average_salary":
                result = query_average_salary(params["city"], params["job_name"], data)
                chart_type = 'bar'
                title = f"{params['city']}{params['job_name']}平均薪资"
                # 确保result是字典格式
                if isinstance(result, (float, int)):
                    result = {title: result}
                chart_type = 'bar'
                title = f"{params['city']}{params['job_name']}平均薪资"  # 添加标题
                
            elif func_name == "query_job_demand":
                result = query_job_demand(params["city"], params["job_name"], data)
                # 确保result是字典格式
                if isinstance(result, int):
                    result = {f"{params['city']}{params['job_name']}岗位数量": result}
                chart_type = 'bar'
                title = f"{params['city']}{params['job_name']}岗位数量"
                # 确保result是字典格式
                if isinstance(result, int):
                    result = {title: result}
                chart_type = 'bar'  # 修改为柱状图
                title = f"{params['city']}{params['job_name']}岗位数量"  # 添加标题
                
            elif func_name == "query_max_salary_job":
                result = query_max_salary_job(params["city"], data, top_n=1)
                chart_type = 'bar'  # 修改为柱状图
                title = f"{params['city']}薪资最高岗位"  # 添加标题
            
            elif func_name == "query_language_related_jobs":
                language = params["language"]
                result = query_language_related_jobs(data, language=language, top_n=params.get("top_n",5))
                chart_type = 'pie'  # 新增：显式设置为饼图
                title = f"{language}相关热门岗位分布"  # 可选：添加标题（若需要）
            
            # 新增：处理学习推荐工具
            elif func_name == "query_cpp_learning_recommendations":
                result = query_cpp_learning_recommendations(data, params.get("top_n", 10))
                chart_type = 'ring'  # 关键修改：将饼图改为环形图
                title = None  # 保持标题不显示
            
            elif func_name == "query_language_jobs":
                language = params["language"]
                result = query_language_jobs(data, language=language, top_n=params.get("top_n",10))
                
                # 保存表格数据
                st.session_state.last_table_data = result if not result.empty else None
                
                # 直接返回查询结果，用于后续在主界面展示表格
                return result

            elif func_name == "query_top_jobs":
                city = params.get("city")
                result = query_top_jobs(city, data, params.get("top_n", 5))
                chart_type = 'bar'
                title = f"{city}地区热门岗位分布" if city else "热门岗位分布"  # 动态城市标题

            # 生成自然语言回答
            prompt = f"""根据以下信息回答问题：
            用户提问：{question}
            查询结果：{result}
            请用简洁自然的语言回答，不超过200字"""
            
            final_response = client.chat.completions.create(
                model="glm-4",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.3
            )
            answer = final_response.choices[0].message.content
            
            # 关键修改：调整保存条件（允许title为None）
            if chart_type:  # 只要chart_type有效就保存，不再强制要求title
                st.session_state.last_chart_data = (result, chart_type, title)  # 保留title（可能为None）
            
            return answer
        else:
            return response.choices[0].message.content
    except Exception as e:
        return f"处理问题时发生错误: {str(e)}"

# Streamlit应用
def main():
    st.set_page_config(page_title="招聘数据智能助手", layout="wide")
    st.title("招聘数据智能助手")
    
    # 统一初始化所有会话状态变量
    if 'messages' not in st.session_state:
        st.session_state.messages = [{"role": "assistant", "content": "您好，我是招聘数据分析助手，请问有什么可以帮助您？"}]
    if 'api_key' not in st.session_state:
        st.session_state.api_key = ""
    if 'last_chart_data' not in st.session_state:  # 改为存储数据而非widget
        st.session_state.last_chart_data = None  # 初始化图表数据状态
    if 'last_table_data' not in st.session_state:  # 新增：初始化表格数据状态
        st.session_state.last_table_data = None  # 初始化为None
    # 添加CSS样式美化表格
    st.markdown("""
    <style>
        .dataframe th {
            background-color: #f0f2f6 !important;
            color: #2c3e50 !important;
            font-weight: 600 !important;
        }
        .dataframe td {
            padding: 10px !important;
        }
        .stDataFrame {
            border-radius: 8px;
            box-shadow: 0 2px 6px rgba(0,0,0,0.1);
            margin-top: 20px;
        }
    </style>
    """, unsafe_allow_html=True)
    # 清空聊天记录的回调函数
    def clear_chat_history():
        st.session_state.messages = [{"role": "assistant", "content": "您好，我是招聘数据分析助手，请问有什么可以帮助您？"}]
        st.session_state.last_chart_data = None  # 清空图表数据
    
    data = load_data()
    
    # 左侧边栏内容
    with st.sidebar:
        st.header("功能选项")
        api_key = st.text_input("请输入您的API Key", 
                              type="password",
                              value=st.session_state.api_key,
                              help="请输入有效的智谱AI API Key")
        if api_key != st.session_state.api_key:  # 只有变化时才更新
            st.session_state.api_key = api_key
            st.rerun()  # 触发重新验证
        st.button('清空聊天记录', on_click=clear_chat_history)
    
    # 显示历史对话
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])
    
    # 显示图表（添加长度检查）
    if st.session_state.last_chart_data:
        # 检查元组长度是否为3，避免解包错误
        if len(st.session_state.last_chart_data) == 3:
            result, chart_type, title = st.session_state.last_chart_data  # 读取title
            plot_top_jobs(result, chart_type=chart_type, title=title)  # 传递title参数
        else:
            st.warning("图表数据格式错误，无法显示")  # 新增：错误提示
    

    # 在主函数main中，展示表格的部分
    if st.session_state.last_table_data is not None:
        st.markdown("### 📋 岗位详细信息")
        st.dataframe(
            st.session_state.last_table_data,
            height=500,
            use_container_width=True,
            hide_index=True,
            column_config={
                "职位名称": "职位名称",
                "公司名称": st.column_config.TextColumn(
                    "公司名称", width="large"
                ),
                "城市": st.column_config.TextColumn(
                    "城市", width="small"
                ),
                "薪资范围": st.column_config.ProgressColumn(
                    "薪资水平",
                    help="根据薪资中位数计算的相对水平",
                    format="%.2f",
                    min_value=0,
                    max_value=100
                )
            }
        )
    elif st.session_state.last_table_data is not None and st.session_state.last_table_data.empty:
        st.warning("未找到相关岗位信息")
    # 只有输入了有效的API key才能使用聊天功能
    if st.session_state.api_key:
        try:
            client = init_ai(st.session_state.api_key)
            if prompt := st.chat_input("请输入您的问题"):
                # 添加用户消息
                st.session_state.messages.append({"role": "user", "content": prompt})
                with st.chat_message("user"):
                    st.write(prompt)
                
                # 处理问题
                with st.spinner("正在思考..."):
                    with st.chat_message("assistant"):
                        answer = process_user_question(prompt, data, st.session_state.api_key)
                        if answer:
                            st.write(answer)
                    st.session_state.messages.append({"role": "assistant", "content": answer})
                
                # 关键修改：在问题处理完成后显示图表（确保last_chart_data已更新）
                if st.session_state.last_chart_data:
                    result, chart_type, title = st.session_state.last_chart_data
                    plot_top_jobs(result, chart_type=chart_type, title=title)
        except Exception as e:
            st.error(f"发生错误: {str(e)}")
            st.session_state.api_key = ""
    else:
        st.warning("请输入有效的API Key以启用聊天功能")

if __name__ == "__main__":
    main()
